package com.googlecode.connectlet.crypto.keystore;

import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.security.Key;
import java.security.KeyStoreException;
import java.security.KeyStoreSpi;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.util.Date;
import java.util.Enumeration;
import java.util.Hashtable;

import com.googlecode.connectlet.crypto.random.RC4PrngCipher;
import com.googlecode.connectlet.crypto.stream.RC4Cipher;
import com.googlecode.connectlet.util.Bytes;
import com.googlecode.connectlet.util.Streams;

class Entry implements Serializable {
	private static final long serialVersionUID = 1L;

	Certificate cert = null;
	byte[] key = null;
	Certificate[] chain = null;
	Date date;

	Entry(Certificate cert) {
		this.cert = cert;
		date = new Date();
	}

	Entry(byte[] key, Certificate[] certChain) {
		this.key = key;
		this.chain = certChain;
		date = new Date();
	}
}

class Entries extends Hashtable<String, Entry> {
	private static final long serialVersionUID = 1L;
}

public abstract class EasyKeyStore extends KeyStoreSpi {
	protected abstract PrivateKey generatePrivate(byte[] encoded);

	private Entries entries;

	private static final int DEFAULT_ITERATION = 1000000;

	private static RC4PrngCipher random = RC4PrngCipher.getInstance();

	@Override
	public Enumeration<String> engineAliases() {
		return entries.keys();
	}

	@Override
	public boolean engineContainsAlias(String alias) {
		return entries.containsKey(alias);
	}

	@Override
	public void engineDeleteEntry(String alias) throws KeyStoreException {
		entries.remove(alias);
	}

	@Override
	public int engineSize() {
		return entries.size();
	}

	@Override
	public Date engineGetCreationDate(String alias) {
		return entries.get(alias).date;
	}

	@Override
	public Certificate engineGetCertificate(String alias) {
		return entries.get(alias).cert;
	}

	@Override
	public String engineGetCertificateAlias(Certificate cert) {
		Enumeration<String> e = entries.keys();
		while (e.hasMoreElements()) {
			String alias = e.nextElement();
			Certificate certKeyStore = entries.get(alias).cert;
			if (certKeyStore != null && certKeyStore.equals(cert)) {
				return alias;
			}
			Certificate[] certs = entries.get(alias).chain;
			if (certs != null && certs[1].equals(cert)) {
				return alias;
			}
		}
		return null;
	}

	@Override
	public Certificate[] engineGetCertificateChain(String alias) {
		return entries.get(alias).chain;
	}

	@Override
	public Key engineGetKey(String alias, char[] password) {
		byte[] key = entries.get(alias).key;
		byte[] pbkey = RC4PBKDF.genKey(32, new String(password),
				Bytes.left(key, 16), DEFAULT_ITERATION);
		return generatePrivate(new RC4Cipher(pbkey).update(Bytes.sub(key, 16)));
	}

	@Override
	public boolean engineIsCertificateEntry(String alias) {
		return entries.get(alias).cert != null;
	}

	@Override
	public boolean engineIsKeyEntry(String alias) {
		return entries.get(alias).key != null;
	}

	@Override
	public void engineSetCertificateEntry(String alias, Certificate cert) {
		entries.put(alias, new Entry(cert));
	}

	@Override
	public void engineSetKeyEntry(String alias, byte[] key, Certificate[] chain) {
		entries.put(alias, new Entry(key, chain));
	}

	@Override
	public void engineSetKeyEntry(String alias, Key key, char[] password, Certificate[] chain) {
		byte[] salt = random.output(16);
		byte[] pbkey = RC4PBKDF.genKey(32, new String(password), salt, DEFAULT_ITERATION);
		byte[] encodedKey = Bytes.add(salt, new RC4Cipher(pbkey).update(key.getEncoded()));
		engineSetKeyEntry(alias, encodedKey, chain);
	}

	@Override
	public void engineLoad(InputStream stream, char[] password) throws IOException {
		if (stream == null) {
			entries = new Entries();
		} else {
			try {
				byte[] pbkey = RC4PBKDF.genKey(32, new String(password),
						Streams.read(stream, 16), DEFAULT_ITERATION);
				InputStream in = new RC4Cipher(pbkey).decrypt(stream);
				entries = (Entries) new ObjectInputStream(in).readObject();
			} catch (Exception e) {
				throw new IOException(e);
			}
		}
	}

	@Override
	public void engineStore(OutputStream stream, char[] password) throws IOException {
		byte[] salt = random.output(16);
		stream.write(salt);
		byte[] pbkey = RC4PBKDF.genKey(32, new String(password), salt, DEFAULT_ITERATION);
		new ObjectOutputStream(new RC4Cipher(pbkey).encrypt(stream)).writeObject(entries);
	}
}